#include <task/task.h>
#include <task/sched.h>

#include <mm/paging.h>
#include <mm/page_allocator.h>
#include <mm/heap.h>

#include <arch/intr.h>
#include <arch/processor.h>
#include <arch/lapic.h>

#include <utils/memutils.h>
#include <utils/stack_emulator.h>

#include <drv/serial.h>

static uint64_t task_alloc_tid() {
	static uint64_t tid_counter = 0;

	return tid_counter++;
}

task_block_t *task_create(void (*entry)(void), uint64_t *page_table, uint32_t flags) {
	if(!page_table) {
		return NULL;
	}

	task_block_t *task = mm_heap_zeroalloc(sizeof(task_block_t));

	task->tid = task_alloc_tid();
	task->task_flags = flags;

	task->kstack_base = mm_pallocator_alloc4k(TASK_KSTACK_SIZE >> PAGE_4KSHIFT);
	task->kstack_top = task->kstack_base + TASK_KSTACK_SIZE;

	mm_paging_map_memory(page_table, task->kstack_base, task->kstack_base, TASK_KSTACK_SIZE >> PAGE_4KSHIFT, PAGE_P | PAGE_RW);

	if(flags & TASK_FLAG_USER) {
		task->ustack_base = mm_pallocator_alloc4k(TASK_USTACK_SIZE >> PAGE_4KSHIFT);
		task->ustack_top = task->ustack_base + TASK_USTACK_SIZE;

		mm_paging_map_memory(kernel_pagetable, task->ustack_base, task->ustack_base, TASK_USTACK_SIZE >> PAGE_4KSHIFT, PAGE_US | PAGE_P | PAGE_RW);
		mm_paging_map_memory(page_table, task->ustack_base, task->ustack_base, TASK_USTACK_SIZE >> PAGE_4KSHIFT, PAGE_US | PAGE_P | PAGE_RW);

		task->task_stack_top = task->ustack_top;
	} else {
		task->task_stack_top = task->kstack_top;
	}

	stackemu_context_t stackemu;
	stackemu_construct(&stackemu, task->task_stack_top);

	intr_regs_frame_t *frame = (intr_regs_frame_t *)stackemu_alloc(&stackemu, sizeof(*frame));
	memset(frame, 0, sizeof(intr_regs_frame_t));

	task->task_stack_top = stackemu.bp;

	frame->cs = flags & TASK_FLAG_USER ? GDT_MAKE_SELECTOR(GDT_USER_CS_IDX, 3) : GDT_MAKE_SELECTOR(GDT_KERNEL_CS_IDX, 0);
	frame->cr3 = (uint64_t)page_table;
	frame->rip = (uint64_t)entry;
	frame->rsp = (uint64_t)task->task_stack_top;
	frame->rflags = 0x202;	// IF = 1, IOPL = 0

	if(flags & TASK_FLAG_USER) {
		frame->ss = GDT_MAKE_SELECTOR(GDT_USER_DS_IDX, 3);
	}

	serial_printf("task_create[%p]: rsp = %p, rip = %p, cr3 = %p\n", task, frame->rsp, frame->rip, frame->cr3);
	serial_printf("task_create[%p]: kstack_base = %p, kstack_top = %p\n", task, task->kstack_base, task->kstack_top);
	serial_printf("task_create[%p]: ustack_base = %p, ustack_top = %p\n", task, task->ustack_base, task->ustack_top);
	serial_printf("task_create[%p]: fxsave_buf = %p, task_stack_top = %p\n", task, frame->fxsave_buf, task->task_stack_top);
	serial_printf("task_create[%p]: &rflags = %p\n", task, &frame->rflags);

	task->task_state = TASK_STATE_READY;

	return task;
}

void task_idle(void) {
	// __early_panic(0x11233, 9);

	while(1) {
		serial_printf("idle task on core %d\n", lapic_get_self_id());
	}
}

task_block_t *task_create_idle() {
	return task_create(task_idle, kernel_pagetable, 0);
}

void task1(void) {
	while(1) {
		serial_printf("task1 on core %d\n", lapic_get_self_id());
	}
}

void task2(void) {
	while(1) {
		serial_printf("task2 on core %d\n", lapic_get_self_id());
	}
}

void task3(void) {
	while(1) {
		serial_printf("task3 on core %d\n", lapic_get_self_id());
	}
}

void task4(void) {
	while(1) {
		serial_printf("task4 on core %d\n", lapic_get_self_id());
	}
}

void task_add_test_tasks() {
	uint64_t *task_pt = (uint64_t *)mm_pallocator_alloc4k(1);
	memset(task_pt, 0, PAGE_4K);

	mm_paging_map_physical_mem(task_pt);

	mm_paging_map_memory_4k(task_pt, (uint64_t)task1 & ~PAGE_4KMASK, (uint64_t)task1 & ~PAGE_4KMASK, PAGE_US | PAGE_P | PAGE_RW);
	mm_paging_map_memory_4k(task_pt, (uint64_t)task2 & ~PAGE_4KMASK, (uint64_t)task2 & ~PAGE_4KMASK, PAGE_US | PAGE_P | PAGE_RW);
	mm_paging_map_memory_4k(task_pt, (uint64_t)task3 & ~PAGE_4KMASK, (uint64_t)task3 & ~PAGE_4KMASK, PAGE_US | PAGE_P | PAGE_RW);
	mm_paging_map_memory_4k(task_pt, (uint64_t)task4 & ~PAGE_4KMASK, (uint64_t)task4 & ~PAGE_4KMASK, PAGE_US | PAGE_P | PAGE_RW);

	sched_add_task(task_create(task1, task_pt, 0));
	sched_add_task(task_create(task2, task_pt, 0));
	sched_add_task(task_create(task3, task_pt, 0));
	sched_add_task(task_create(task4, task_pt, 0));
}